/*
** Copyright (C) 2010 Calin Radoni
**
** This file is part of CHScanner.
**
** CHScanner is free software: you can redistribute it and/or modify
** it under the terms of the GNU General Public License as published by
** the Free Software Foundation, either version 3 of the License, or
** (at your option) any later version.
**
** CHScanner is distributed in the hope that it will be useful,
** but WITHOUT ANY WARRANTY; without even the implied warranty of
** MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
** GNU General Public License for more details.
**
** You should have received a copy of the GNU General Public License
** along with CHScanner. If not, see <http://www.gnu.org/licenses/>.
*/

#include "StdAfx.h"
#include "ScanMethodIgmp.h"

//--------------------------------------------------------------------

ScanMethodIgmp::ScanMethodIgmp(void) : ScanMethod()
{
	myType = __wmtIgmp;
}

ScanMethodIgmp::~ScanMethodIgmp(void)
{
}

//--------------------------------------------------------------------
int ScanMethodIgmp::InternalInit(void)
{
	CString str, ssv, ssg, ssd;
	MAddressIP taip;
	int pos;

	// version grpIpAddress dstIpAddress

	if(workPhase!=__wptDiscovery) return MSR_ERR_PHASE;

	FillRandomBuffer();

	str = scanParameters.Trim();
	if(str.IsEmpty()){
		// no parameters supplied
		queryVersion = 2;
		grpIpAddress.S_un.S_addr = 0;
		dstIpAddress.S_un.S_un_b.s_b1 = 224;
		dstIpAddress.S_un.S_un_b.s_b2 = 0;
		dstIpAddress.S_un.S_un_b.s_b3 = 0;
		dstIpAddress.S_un.S_un_b.s_b4 = 1;
		return MSR_ERR_OK;
	}

	pos = str.Find(' ');
	if(pos==-1){
		// just the version
		queryVersion = (BYTE)_tstoi(str);
		if(queryVersion<1 || queryVersion>3)
			return MSR_ERR_PARAMETER;

		grpIpAddress.S_un.S_addr = 0;

		dstIpAddress.S_un.S_un_b.s_b1 = 224;
		dstIpAddress.S_un.S_un_b.s_b2 = 0;
		dstIpAddress.S_un.S_un_b.s_b3 = 0;
		dstIpAddress.S_un.S_un_b.s_b4 = 1;
		return MSR_ERR_OK;
	}

	// version and others
	ssv = str.Left(pos);
	str = str.Mid(pos+1);
	queryVersion = (BYTE)_tstoi(ssv);
	if(queryVersion<1 || queryVersion>3)
		return MSR_ERR_PARAMETER;

	str = str.Trim();
	pos = str.Find(' ');
	if(pos==-1){
		// group address only
		ssg = str;
		ssd = _T("224.0.0.1");
	}
	else{
		ssg = str.Left(pos);
		ssd = str.Mid(pos+1);
		ssd = ssd.Trim();
	}

	// group address validation (0 or multicast IPv4 address)
	taip.Set(ssg);
	str = taip.GetS();
	if(ssg.CompareNoCase(str)!=0)
		return MSR_ERR_PARAMETER;
	grpIpAddress.S_un.S_addr = taip.GetD();
	if(grpIpAddress.S_un.S_addr!=0){
		if(grpIpAddress.S_un.S_un_b.s_b1<224 || grpIpAddress.S_un.S_un_b.s_b1>239)
			return MSR_ERR_PARAMETER;
	}

	// destination address validation (multicast IPv4 address only)
	taip.Set(ssd);
	dstIpAddress.S_un.S_addr = taip.GetD();
	if(dstIpAddress.S_un.S_un_b.s_b1<224 || dstIpAddress.S_un.S_un_b.s_b1>239)
		return MSR_ERR_PARAMETER;

	return MSR_ERR_OK;
}

//--------------------------------------------------------------------

int ScanMethodIgmp::InternalExecute(void)
{
	ScanPacket* pack;
	bool res;

	// do NOT set igmpMaxResponseTime here !!!
	while(PacketsUsed()>0){
		CheckResponses();
		CheckDeaths();
		SleepEx(1, TRUE);
	}

	// set igmpMaxResponseTime here !!! 
	igmpMaxResponseTime = packetBuilderEx->GetIgmpMaxResponseTime();
	if(igmpMaxResponseTime==0){
		igmpMaxResponseTime = 100;	// RFC specified for v1 compatibility
	}
	else{
		if(igmpMaxResponseTime>128){
			// RFC 3376
			BYTE mant;
			BYTE exp;

			mant  = igmpMaxResponseTime & 0x0F;
			exp   = igmpMaxResponseTime >> 4;
			exp  &= 0x07;

			igmpMaxResponseTime = (mant | 0x10) << (exp + 3);
		}
	}

	pack = FindFirstUnusedPacket();
	while(pack==NULL){
		SleepEx(waitTime, TRUE);
		CheckResponses();
		CheckDeaths();
		pack = FindFirstUnusedPacket();
	}

	if(pack!=NULL){
		pack->Initialize(true);
		pack->sTime       = GetTickCount();
		pack->addr        = 0;
		pack->port        = 0;
		pack->methodType  = myType;
		pack->fullAddress = NULL;
		pack->machineInfo = NULL;

		res = SendIgmpQuery(pack);
		if(!res)
			pack->Initialize(false);
	}
	else res = false;

	while(PacketsUsed()>0){
		CheckResponses();
		CheckDeaths();
		SleepEx(1, TRUE);
	}

	if(hWnd!=NULL)
		PostMessage(hWnd, CHS_MSG_PACKETS_SENT, (WPARAM)0, (LPARAM)packSent);

	if(!res) return MSR_ERR_FAILED;

	return MSR_ERR_OK;
}

//--------------------------------------------------------------------

void ScanMethodIgmp::ComputeRandomValues(void)
{
	// nothing
}

bool ScanMethodIgmp::SendIgmpQuery(ScanPacket* scanPackIn)
{
	BYTE srcMac[6];
	BYTE grpAddr[4];
	bool res;

	if(scanPackIn==NULL) return false;

	grpAddr[0] = grpIpAddress.S_un.S_un_b.s_b1;
	grpAddr[1] = grpIpAddress.S_un.S_un_b.s_b2;
	grpAddr[2] = grpIpAddress.S_un.S_un_b.s_b3;
	grpAddr[3] = grpIpAddress.S_un.S_un_b.s_b4;
	res = packetBuilderEx->BuildPacketIGMPQuerySpecific(winPcap->GetLocalIP(), dstIpAddress.S_un.S_addr, queryVersion, &grpAddr[0]);

	if(!res) return false;

	winPcap->FillSrcMAC(&srcMac[0]);

	res = packetBuilderEx->AddEthernetHeaderForIPMulticast(srcMac, dstIpAddress.S_un.S_addr, 0x0800);
	if(!res) return false;

	return SendPacket(packetBuilderEx->GetPacket(), (int)packetBuilderEx->GetPacketSize());
}

//--------------------------------------------------------------------

void ScanMethodIgmp::ProcessIgmp(BYTE* dataIn, int sizeIn, BYTE* srcMAC)
{
	ScanPacket        *pack;
	IPHeaderPG        *iph;
	IGMPHeaderPG      *igmph;
	IGMPReportPG      *igmpr;
	IGMPGroupRecordPG *igmpg;
	int               dataOffset;
	MMachineInfo      mach;
	MFullAddress      addr;
	MIGMPPacket       iPack;
	in_addr           inAddr, ga;
	int               i, cnt, gCnt, pos;

	if(sizeIn<32) return;		// sizeof(IPHeaderPG) + options(=4bytes) + 8_bytes_for_min_IGMP

	pack = FindPacket(0, 0, myType);
	if(pack==NULL) return;

	iph = (IPHeaderPG*)dataIn;
	dataOffset  = (int)(iph->HLV & 0x0F);
	dataOffset *= 4;
	if(sizeIn<(dataOffset + 8)) return; // incomplete packet

	igmph = (IGMPHeaderPG*)&dataIn[dataOffset];
	switch(igmph->Type){
		case 0x12 : // Version 1 Membership Report
		case 0x16 : // Version 2 Membership Report
			mach.Initialize();
			iPack.type = igmph->Type;
			ga.S_un.S_un_b.s_b1 = igmph->GroupAddress[0];
			ga.S_un.S_un_b.s_b2 = igmph->GroupAddress[1];
			ga.S_un.S_un_b.s_b3 = igmph->GroupAddress[2];
			ga.S_un.S_un_b.s_b4 = igmph->GroupAddress[3];
			iPack.addr.Set(ga.S_un.S_addr);
			mach.AddIgmpPackNoDuplicates(&iPack);
			break;

		case 0x22 : // Version 3 Membership Report
			igmpr = (IGMPReportPG*)&dataIn[dataOffset];
			cnt   = ntohs(igmpr->NoOfGroupRecords);
			pos   = dataOffset + 8; // where 8 is the size of IGMPReportPG
			mach.Initialize();
			while(cnt>0){
				if(sizeIn<pos) return;

				igmpg = (IGMPGroupRecordPG*)&dataIn[pos];
				gCnt  = ntohs(igmpg->NumberOfSources);
				pos += 4*gCnt;
				if(sizeIn<pos) return;

				iPack.type = igmpr->Type;
				ga.S_un.S_un_b.s_b1 = igmpg->MulticastAddress[0];
				ga.S_un.S_un_b.s_b2 = igmpg->MulticastAddress[1];
				ga.S_un.S_un_b.s_b3 = igmpg->MulticastAddress[2];
				ga.S_un.S_un_b.s_b4 = igmpg->MulticastAddress[3];
				iPack.addr.Set(ga.S_un.S_addr);
				iPack.recordType = igmpg->RecordType;

				mach.AddIgmpPackNoDuplicates(&iPack);
	
				cnt--;
				pos += 8;	// where 8 is the size of IGMPGroupRecordPG
			}
			break;

		default :
			return;
	}

	inAddr.S_un.S_un_b.s_b1 = iph->Source[0];
	inAddr.S_un.S_un_b.s_b2 = iph->Source[1];
	inAddr.S_un.S_un_b.s_b3 = iph->Source[2];
	inAddr.S_un.S_un_b.s_b4 = iph->Source[3];

	// add address srcIP:srcMAC

	addr.Initialize();

	for(i=0; i<6; i++)
		addr.addrMAC.addr[i] = srcMAC[i];

	addr.addrIP.Set(inAddr.S_un.S_addr);

	mach.address.LoadFromPointer(&addr);

	machInfoList->Add(&mach);
}

//--------------------------------------------------------------------
